# 用来存放一些和服务器交互的类
import asyncio
import os
from enum import Enum
from typing import Any

from mcp.client.sse import sse_client
from mcp.client.stdio import stdio_client, StdioServerParameters
from contextlib import AsyncExitStack
from mcp.client.session import ClientSession
from models import MCPFunction, MCPFunctionType
from mcp.types import AnyUrl
from mcp.shared.exceptions import McpError


class MCPTransport(Enum):
    """
    mcp连接方式
    """
    STDIO = "stdio"
    SSE = "sse"


class MCPServer:
    """
    这个类代表的是一个MCP Server，里面封装了和MCP Server的连接、通信等，以及获取Session对象，服务的Tools、Resources 、Prompt
    """

    def __init__(
            self,
            name: str,  # 服务名称
            transport: MCPTransport = MCPTransport.STDIO,
            cmd: str | None = None,
            args: list[str] | None = None,
            env: dict[str, Any] | None = None,
            url: str | None = None,
    ):  # python 中钩子函数不能异步
        self.name = name
        self.transport = transport
        if self.transport == MCPTransport.STDIO:
            assert cmd is not None  # 断言是为假的时候抛异常
            assert args is not None
            self.cmd = cmd
            self.args = args
            self.env = env
        else:
            assert url is not None
            self.url = url
            self.env = env if env else {}

        # 异步上下文堆栈
        self._exit_stack = AsyncExitStack()
        # 与服务器通话的Session对象
        self.session: ClientSession | None = None
        # 保存当前MCP Server所有的Function（Tool、Resource、Resource Template、Prompt）
        self.functions: dict[str, MCPFunction] = {}

    async def initialize(self):
        # 初始化操作，连接到MCP Server，以及获取Session对象
        if self.transport == MCPTransport.STDIO:
            params = StdioServerParameters(
                command=self.cmd,
                args=self.args,
                env=self.env, # 从mcp.json中读取的"env"
            )
            reader_stream, writer_stream = await self._exit_stack.enter_async_context(stdio_client(params))
            self.session = await self._exit_stack.enter_async_context(ClientSession(reader_stream, writer_stream))
        else:
            # 从mcp.json的env中获取Authorization
            auth = self.env.get("Authorization") if self.env else ""
            reader_stream, writer_stream = await self._exit_stack.enter_async_context(
                # url: https://127.0.0.1/8000/sse
                sse_client(self.url, headers={"Authorization": f"Bearer {auth}"})
            )
            self.session = await self._exit_stack.enter_async_context(ClientSession(reader_stream, writer_stream))
        # session初始化
        await self.session.initialize()
        # 获取MCP Server的所有Tools、Resources 、Prompt
        await self.fetch_functions()

    async def fetch_functions(self):
        """获取各种Function(Tools、Resources 、Prompt等等)"""
        assert self.session is not None
        # 1.获取Tools
        tools = await self.fetch_tools()
        for tool in tools:
            tool_name = tool.name.replace(" ", "_")
            self.functions[tool_name] = MCPFunction(
                type_=MCPFunctionType.TOOL,
                name=tool_name,
                origin_name=tool.name,
                server_name=self.name,
                description=tool.description,
                input_schema=tool.inputSchema,
            )
        # 2.获取resource
        resources = await self.fetch_resources()
        for resource in resources:
            resource_name = resource.name.replace(" ", "")
            self.functions[resource_name] = MCPFunction(
                type_=MCPFunctionType.RESOURCE,
                name=resource_name,
                origin_name=resource.name,
                server_name=self.name,
                description=resource.description,
                # 如果是resource类型，则uri为AnyUrl类型
                uri=resource.uri,
            )
        # 3.获取resourceTemplate
        resource_templates = await self.fetch_resource_template()
        for resource_template in resource_templates:
            resource_template_name = resource_template.name.replace(" ", "_")
            self.functions[resource_template_name] = MCPFunction(
                type_=MCPFunctionType.RESOURCE_TEMPLATE,
                name=resource_template_name,
                origin_name=resource_template.name,
                server_name=self.name,
                description=resource_template.description,
                # 如果是resource_template类型，则uri为str类型
                uri=resource_template.uriTemplate,
            )
        # 4.获取prompt
        prompts = await self.fetch_prompts()
        for prompt in prompts:
            prompt_name = prompt.name.replace(" ", "_")
            self.functions[prompt_name] = MCPFunction(
                type_=MCPFunctionType.PROMPT,
                name=prompt_name,
                origin_name=prompt.name,
                server_name=self.name,
                description=prompt.description,
                arguments=prompt.arguments,
            )

    async def call_function(self, function_name: str, arguments: dict[str, Any] | None = None):
        """调用Function Call"""
        assert self.session is not None
        function = self.functions[function_name]
        function_type = function.type_
        if function_type == MCPFunctionType.TOOL:
            response = await self.session.call_tool(name=function.origin_name, arguments=arguments)
            return response.content[0].text
        elif function_type == MCPFunctionType.RESOURCE:
            response = await self.session.read_resource(uri=function.uri)
            return response.contents[0].text
        elif function_type == MCPFunctionType.RESOURCE_TEMPLATE:
            # resource_template类型将参数格式化到uri中
            # uri: file://{filename}.format(filename=xxx)
            uri = AnyUrl(function.uri.format(**arguments))
            response = await self.session.read_resource(uri)
            return response.contents[0].text
        else:
            response = await self.session.get_prompt(name=function.origin_name, arguments=arguments)
            return response.content.text

    async def fetch_tools(self):
        assert self.session is not None
        try:
            return (await self.session.list_tools()).tools
        except McpError:
            return []

    async def fetch_resources(self):
        assert self.session is not None
        try:
            return (await self.session.list_resources()).resources
        except McpError:
            return []

    async def fetch_resource_template(self):
        assert self.session is not None
        try:
            return (await self.session.list_resource_templates()).resourceTemplates
        except McpError:
            return []

    async def fetch_prompts(self):
        assert self.session is not None
        try:
            return (await self.session.list_prompts()).prompts
        except McpError:
            return []

    async def aclose(self):
        await self._exit_stack.aclose()

    async def __aenter__(self):
        # 进入到异步上下文当中做什么
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        # 退出异步上下文做什么
        await self.aclose()


class MCPServerManager:
    """
    由于客户端能同时连接多个MCP Server，因此我们可以单独定义一个Manage类用来管理这些MCP Server
    """

    def __init__(self, mcp_server_dicts: dict):
        self.mcp_server_dicts = mcp_server_dicts
        # self.servers: list[MCPServer] = []
        self.servers: dict[str, MCPServer] = {}
        self._exit_stack: AsyncExitStack = AsyncExitStack()
        self.all_functions: dict[str, MCPFunction] = {}  # 所有的MCP Server的Functioon总和

    async def initialize(self):
        for name, mcp_server_dict in self.mcp_server_dicts.items():
            # 1.创建好所有的MCP Server对象
            transport = MCPTransport.STDIO
            if not mcp_server_dict.get("command"):
                transport = MCPTransport.SSE
            server = await self._exit_stack.enter_async_context(
                MCPServer(
                    name=name,
                    transport=transport,
                    cmd=mcp_server_dict.get("command"),
                    args=mcp_server_dict.get('args'),
                    env=mcp_server_dict.get('env'),
                    url=mcp_server_dict.get('url'),
                )
            )
            self.servers[name] = server
            # 2.通过server可以获取所有MCP Server的函数并保存起来，也存储成字典形式
            self.all_functions.update(server.functions) # update添加到字典中

    async def call_function(self, name: str, arguments: dict[str, Any] | None = None):
        function = self.all_functions[name]
        server = self.servers[function.server_name]
        return await server.call_function(name, arguments)

    async def aclose(self):
        await self._exit_stack.aclose()

    async def __aenter__(self):
        # 进入到异步上下文当中做什么
        await self.initialize()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        # 退出异步上下文做什么
        await self.aclose()


async def main1():
    server = MCPServer(
        name="filesystem",
        cmd="npx",
        args=["-y", "@modelcontextprotocol/server-filesystem", "D:/githubProject/pythonMCPstudy/data"]
    )
    await server.initialize()
    result = await server.call_function("read_file", {"path": "D:/githubProject/pythonMCPstudy/data/yangshen.txt"})
    print(result)
    await server.aclose()  # 手动关闭有点麻烦


async def main2():
    """
    本调用方式是借助异步事件自动实现初始化和关闭
    :return:
    """
    async with MCPServer(
            name="filesystem",
            cmd="npx",
            args=["-y", "@modelcontextprotocol/server-filesystem", "D:/githubProject/pythonMCPstudy/data"]
    ) as server:
        result = await server.call_function("read_file", {"path": "D:/githubProject/pythonMCPstudy/data/yangshen.txt"})
        print(result)


async def main():
    """
    本调用方式是借助异步上下问堆栈自动实现初始化和关闭，
    :return:
    """
    exit_stack = AsyncExitStack()
    server = await exit_stack.enter_async_context(MCPServer(
        name="sqlite",
        cmd="uvx",
        args=["mcp-server-sqlite", "--db-path", "D:/githubProject/pythonMCPstudy/data/db.sqlite"]
    ))
    result = await server.call_function(function_name="Business Insights Memo")
    print(result)
    await exit_stack.aclose()


if __name__ == '__main__':
    """调用方式用main2/main都可以"""
    asyncio.run(main())
